import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from linear_regression import LinearRegression

data = pd.read_csv('data/lonlat.csv')
# data = pd.read_csv('data/国家幸福GDP.csv')

train_data = data.sample(frac=0.8) # 取数据中的80%作为训练数据
test_data = data.drop(train_data.index) # 取剩下的数据为测试数据

input_param_name = 'time'  # 输入列名 x
output_param_name = 'lon' # 输出列名 y

x_train = train_data[[input_param_name]].values # labels
y_train = train_data[[output_param_name]].values
x1 = train_data[['time']].values

x_test = test_data[[input_param_name]].values
y_test = test_data[[output_param_name]].values
x2 = test_data[['time']].values


plt.scatter(x_train, y_train, color='red',label='Train data')
plt.scatter(x_test, y_test, color='blue',label='Test data')
plt.xlabel(input_param_name)
plt.ylabel(output_param_name)
plt.title(output_param_name)
plt.legend()
plt.show()

num_iterations = 500 # 迭代次数
learning_rate = 0.01  # 学习率
linear_regression=LinearRegression(x_train,y_train)
(theta,cost_history)=linear_regression.train(learning_rate,num_iterations)

print('开始时的损失：',cost_history[0])
print('训练后的损失：',cost_history[-1])

plt.plot(range(num_iterations),cost_history)
plt.xlabel('iterations')
plt.ylabel('cost')
plt.show() # 显示梯度下降

# 测试
predictions_num = 100
x_predictions = np.linspace(x_train.min(), x_train.max(), predictions_num).reshape(predictions_num,1)
y_predictions = linear_regression.predict(x_predictions)

plt.scatter(x_train, y_train, color='red',label='Train data')
plt.scatter(x_test, y_test, color='blue',label='Test data')
plt.plot(x_predictions, y_predictions,label='Linear Regression')
plt.xlabel('x')
plt.ylabel('y')
plt.title(output_param_name)
plt.legend()
plt.show()